import numpy as np

# 状态空间模型参数
A = np.array([[1.0, 0.0],
              [0.0, 1.0]])  # 系统动态矩阵
B = np.array([[0.2,0.6],
              [0.8,0.4]])  # 输入矩阵
C = np.eye(2)  # 输出矩阵
D = np.zeros((2, 1))  # 直通矩阵

# MPC 参数
N = 10  # 预测时域
Q = np.eye(2)  # 状态权重
R = np.eye(1)  # 控制权重

# 目标状态
r = np.array([[5],
              [3]])

# 初始化状态
x = np.array([[0],
              [0]])


def mpc_control(A, B, Q, R, N, x, r):
    n_states = A.shape[0]
    n_controls = B.shape[1]

    # 构建优化问题
    H = np.zeros((N * n_controls, N * n_controls))
    F = np.zeros((N * n_controls, 1))

    for i in range(N):
        H[i * n_controls:(i + 1) * n_controls, i * n_controls:(i + 1) * n_controls] = R
        F[i * n_controls:(i + 1) * n_controls, 0] = 0

    G = np.zeros((N * n_states, N * n_controls))
    E = np.zeros((N * n_states, n_states))
    L = np.zeros((N * n_states, 1))

    for i in range(N):
        for j in range(i + 1):
            G[i * n_states:(i + 1) * n_states, j * n_controls:(j + 1) * n_controls] = np.linalg.matrix_power(A,
                                                                                                             i - j) @ B
        E[i * n_states:(i + 1) * n_states, :] = np.linalg.matrix_power(A, i + 1)
        L[i * n_states:(i + 1) * n_states, :] = np.linalg.matrix_power(A, i + 1) @ x - r

    Q_bar = np.kron(np.eye(N), Q)
    H += G.T @ Q_bar @ G
    F += (G.T @ Q_bar @ L).reshape(-1, 1)

    # 计算最优控制输入
    u_opt = -np.linalg.inv(H) @ F

    # 取第一个控制输入
    u = u_opt[:n_controls]

    return u


# 仿真时间
T = 20
x_history = []
u_history = []

for t in range(T):
    u = mpc_control(A, B, Q, R, N, x, r)
    x = A @ x + B @ u
    print(u)
    print(x)

    x_history.append(x.flatten())
    u_history.append(u.flatten())

x_history = np.array(x_history)
u_history = np.array(u_history)

print("状态轨迹：\n", x_history)
print("控制输入：\n", u_history)
